#include <assert.h>
#include <clock/clock.h>
#include <cpio/cpio.h>
#include <cspace/cspace.h>
#include <elf/elf.h>
#include <fcntl.h>
#include <string.h>
#include <strings.h>

#include "../bitset.h"
#include "../elf.h"
#include "../frametable.h"
#include "../pagetable.h"
#include "../process.h"
#include "../task.h"
#include "../ut_manager/ut.h"
#include "../vmem_layout.h"

#include "../debug.h"
#define verbose SOS_DEBUG_PROC_CREATE
#include <sys/debug.h>
#include <sys/panic.h>

#define SOS_PROCESS_IPC_ATTRS   (page_rw)
#define SOS_PROCESS_STACK_ATTRS (page_rw)
#define SOS_PROCESS_HEAP_ATTRS  (page_rwx)

/* State required to resume the creation of a process asynchronously */
typedef struct process_create_data {
    sos_pcb *pcb;
    char *elf_base;
    /* Info to call back calling process */
    process_create_cb cb;
    void *cookie;
} process_create_data;

/* Pointer to the cpio boot image */
extern char _cpio_archive[];

/* Global SOS fault endpoint */
extern seL4_CPtr _sos_ipc_ep_cap;

/* Global vfs */
extern sos_vfs *vfs;

/* The SOS process table - managed by libut */
sos_pcb *_process_table;
int _num_procs;

/* Most recent PID */
int _last_pid;

/* Bitset for PIDs */
bitset *pids;

/* Forward declarations of various static functions */
static int _set_pid(process_create_data *data);
static int _create_sos_pd(process_create_data *proc);
static void _create_sos_pd_cb(int err, void *data);
static int _create_vspace(sos_pcb *pcb);
static int _create_cspace(sos_pcb *pcb);
static int _create_ipc_buffer(process_create_data *proc);
static void _enable_ipc(int err, seL4_Word sos_vaddr, void *cookie);
static int _create_guard_page(sos_pcb *pcb, seL4_Word proc_vaddr);
static int _create_stack_frame(sos_pcb *pcb);
static int _create_heap(sos_pcb *pcb);
static int _create_tcb(sos_pcb *pcb, seL4_Word priority);
static int _load_elf(sos_pcb *pcb, char* app_name, process_create_data *data);
static int _create_stderrout(sos_pcb *pcb);
static void _start_process(int err, seL4_Word pc, void *cookie);

/* Perform callback after process creation is complete */
static void _reply(int err, int pid, process_create_cb cb, void *cookie) {
    dprintf(5, "process_create_reply: calling back for PID %d\n", pid);
    (*cb)(err, pid, cookie);
}

static void _fail(int err, process_create_data *proc) {
    _reply(err, -1, proc->cb, proc->cookie);
    if (proc->pcb) {
        proc->pcb->abort = 1;
        process_delete(proc->pcb, NULL, NULL);
    }
    free(proc);
}

/* PID -> PCB */
int lookup_pcb(int pid, sos_pcb **pcb) {
    sos_pcb *res = NULL;
    HASH_FIND_INT(_process_table, &pid, res);
    if (res == NULL) {
        return SOS_PROCESS_NOENT;
    } else {
        (*pcb) = res;
        return SOS_PROCESS_SUCCESS;
    }
}

/* Create a new process from CPIO */
int process_create(char *app_name, process_create_cb cb, void *cookie) {
    int err;

    dprintf(5, "abcde\n");
    dprintf(5, "process_create: %.128s\n", app_name);

    process_create_data *data = malloc(sizeof(struct process_create_data));
    if (data == NULL) {
        dprintf(6, "process_create: OOM\n");
        _reply(SOS_PROCESS_OOM, -1, cb, cookie);
        return SOS_PROCESS_OOM;
    }
    dprintf(8, "bzero procreate data\n");
    bzero(data, sizeof(struct process_create_data));
    dprintf(8, "bzero procreate data done\n");

    data->pcb = malloc(sizeof(struct sos_pcb));
    if (data->pcb == NULL) {
        dprintf(5, "process_create: could not allocate for PCB\n");
        _fail(SOS_PROCESS_OOM, data);
        return SOS_PROCESS_OOM;
    } else {
        dprintf(6, "process_create: PCB is at %p\n", data->pcb);
    }
    bzero(data->pcb, sizeof(struct sos_pcb));

    data->pcb->stime = time_stamp();
    dprintf(6, "\nCopying app name...\n");
    strncpy(data->pcb->app_name, app_name, SOS_PATH_MAX);
    data->cb = cb;
    data->cookie = cookie;

    dprintf(3, "\nStarting \"%s\"...\n", app_name);
    dprintf(5, "\tSetting PID...\n");
    err = _set_pid(data);
    if (err) {
        dprintf(6, "Could not get PID\n");
        _fail(err, data);
        return err;
    }

    dprintf(5, "\tCreating SOS page directory...\n");
    err = _create_sos_pd(data);
    if (err) {
        _fail(err, data);
        return err;
    }
    // The rest happens in CPS
    return SOS_PROCESS_SUCCESS;
}

/* Obtain and set a PID for the new process.
 * O(1) in average case, O(N) degenerate case (more than PID_MAX procs)
 * Uses hardware FFS instruction, so runs in very quick N/32 time */
static int _set_pid(process_create_data *data) {
    // Try to increment _last_pid.
    int err;
    sos_pcb *pcb = data->pcb;
    int cur_pid = _last_pid + 1;
    sos_pcb *taken = NULL;

    if (!pids) {
        dprintf(6, "Allocating PID set...\n");
        bitset_alloc(&pids, SOS_PID_MAX - SOS_PID_MIN);
    }

    dprintf(1, "bs size: %d\n", pids->size);

    if (bitset_test(pids, cur_pid)) {
        dprintf(6, "\tPID %d is already taken.\n", cur_pid);
        // cur_pid is taken
        cur_pid = bitset_ffz(pids);
        dprintf(6, "\tfirst free PID: %d\n", cur_pid);
        if (cur_pid < 0 || cur_pid > pids->size) {
            return SOS_PROCESS_MAXPROC;
        }
    }

    // cur_pid is known to be free
    bitset_set(pids, cur_pid);
    pcb->pid = cur_pid + SOS_PID_MIN;
    _last_pid = cur_pid;
    HASH_ADD_INT(_process_table, pid, pcb);
    return SOS_PROCESS_SUCCESS;
}

static int _create_sos_pd(process_create_data *proc) {
    return sos_pager_create(proc->pcb, &_create_sos_pd_cb, (void *) proc);
}

static void _create_sos_pd_cb(int err, void *data) {
    process_create_data *proc = (process_create_data *) data;

    if (err) {
        _fail(err, proc);
        return;
    }

    dprintf(5, "SOS page directory allocated at %p\n", proc->pcb->sos_pd);

    dprintf(5, "\tCreating vspace...\n");
    err = _create_vspace(proc->pcb);
    if (err) {
        _fail(err, proc);
        return;
    }

    dprintf(5, "\tCreating cspace...\n");
    err = _create_cspace(proc->pcb);
    if (err) {
        _fail(err, proc);
        return;
    }
    dprintf(5, "\tCreating IPC buffer...\n");
    err = _create_ipc_buffer(proc);
    if (err) {
        _fail(err, proc);
        return;
    }
}

/* We defer to ut for kernel objects, for now.
 * SOS kernel objects will never be swappable, so this is OK. */
static int _create_vspace(sos_pcb *pcb) {
    int err;
    pcb->seL4_pd_addr = ut_alloc(seL4_PageDirBits); // 14 bits = 16K
    if (pcb->seL4_pd_addr == 0) {
        dprintf(6, "No memory for new Page Directory\n");
        return SOS_PROCESS_OOM;
    }
    dprintf(7, "create_vspace: got paddr %p for sel4PD\n", pcb->seL4_pd_addr);
    err = cspace_ut_retype_addr(pcb->seL4_pd_addr,
                                seL4_ARM_PageDirectoryObject,
                                seL4_PageDirBits,
                                cur_cspace,
                                &(pcb->seL4_pd));
    if (err) dprintf(6, "Failed to retype memory for PD");
    return err;
}

/* We create a simple Level-1 cspace, because SOS processes get very
   few caps besides their page directory and main endpoints. */
static int _create_cspace(sos_pcb *pcb) {
    pcb->cspace = cspace_create(1);
    if (pcb->cspace == NULL) {
        // Need to abort here! This is real bad
        return SOS_PROCESS_OOM;
    }
    return SOS_PROCESS_SUCCESS;
}

static int _create_ipc_buffer(process_create_data *proc) {
    int err;

    dprintf(7, "process_create_ipc: creating IPC region\n");
    err = sos_create_region(proc->pcb, SOS_PROCESS_IPC_BUFFER, 1,
                            SOS_PROCESS_IPC_ATTRS);
    if (err) return err;

    // Make sure this page is present and unswappable
    dprintf(7, "process_create_ipc: pinning IPC buffer\n");
    err = sos_page_pin(proc->pcb, SOS_PROCESS_IPC_BUFFER,
                       &_enable_ipc, (void *) proc);
    return err;
}

static void _enable_ipc(int err, seL4_Word sos_vaddr, void *cookie) {
    process_create_data *proc = (process_create_data *) cookie;

    if (err) {
        dprintf(6, "process_create_enable_ipc: could not pin IPC buf (%d)\n",
                err);
        _fail(err, proc);
        return;
    }

    // Record the faulted frame address
    proc->pcb->ipc_vaddr = sos_vaddr;

    // Store the process capability for the IPC buffer into our PCB
    dprintf(6, "\t\tCopying process cap...\n");
    err = frame_process_cap(proc->pcb->ipc_vaddr,
                            proc->pcb,
                            &(proc->pcb->ipc_cap));
    // Broken invariant if that failed
    conditional_panic(err, "Failed to obtain IPC process cap");

    // Copy the fault endpoint to the user app to enable IPC
    dprintf(6, "\t\tMinting cap...\n");
    proc->pcb->ep_cap = cspace_mint_cap(proc->pcb->cspace,
                                       cur_cspace,
                                       _sos_ipc_ep_cap,
                                       seL4_AllRights,
                                       seL4_CapData_Badge_new(proc->pcb->pid));

    // Make sure fault endpoint is in the first cspace slot
    if (!(proc->pcb->ep_cap == 1)) {
        _fail(1, proc);
        return;
    }

    // remainder of process creation
    dprintf(5, "\tCreating stack frame...\n");
    err = _create_stack_frame(proc->pcb);
    if (err) {
        _fail(err, proc);
        return;
    }

    dprintf(5, "\tCreating heap...\n");
    err = _create_heap(proc->pcb);
    if (err) {
        _fail(err, proc);
        return;
    }

    dprintf(5, "\tCreating seL4 TCB...\n");
    err = _create_tcb(proc->pcb, SOS_PROCESS_PRIORITY);
    if (err) {
        _fail(err, proc);
        return;
    }

    dprintf(5, "\tLoading ELF...\n");
    err = _load_elf(proc->pcb, proc->pcb->app_name, proc);
    if (err) {
        _fail(err, proc);
        return;
    }

    dprintf(5, "\tCreating children list...\n");
    proc->pcb->children = list_init();
    if (proc->pcb->children == NULL) {
        dprintf(5, "FAILED to create list head\n");
        _fail(SOS_PROCESS_OOM, proc);
        return;
    }
}

/* All our guard pages are handled through the region system */
static int _create_guard_page(sos_pcb *pcb, seL4_Word proc_vaddr) {
    return sos_create_region(pcb, proc_vaddr, 1, page_guard);
}

/* Statically-sized stack frame, for now */
static int _create_stack_frame(sos_pcb *pcb) {
    int err;

    err = sos_create_region(pcb, SOS_PROCESS_STACK_BOTTOM,
              /* Number of frames in the fixed-size stack */
              PAGE_SHIFT(SOS_PROCESS_STACK_TOP - SOS_PROCESS_STACK_BOTTOM),
              SOS_PROCESS_STACK_ATTRS);

    if (err) return err;

    // Allocate stack guard page
    err = _create_guard_page(pcb, SOS_PROCESS_STACK_GUARD);
    return err;
}

/* Statically-sized initial heap, for now */
static int _create_heap(sos_pcb *pcb) {
    int err;
    err = sos_create_region(pcb, SOS_PROCESS_HEAP_START,
              /* Number of pages in the initial heap region */
              PAGE_SHIFT(SOS_PROCESS_BRK - SOS_PROCESS_HEAP_START),
              SOS_PROCESS_HEAP_ATTRS);
    if (err) return err;

    // Allocate heap guard page
    err = _create_guard_page(pcb, SOS_PROCESS_HEAP_GUARD);
    return err;
}


static int _create_tcb(sos_pcb *pcb, seL4_Word priority) {
    int err;
    pcb->tcb_addr = ut_alloc(seL4_TCBBits);

    // Create a new seL4 TCB object
    if (pcb->tcb_addr == 0) return SOS_PROCESS_OOM;

    err =  cspace_ut_retype_addr(pcb->tcb_addr,
                                 seL4_TCBObject,
                                 seL4_TCBBits,
                                 cur_cspace,
                                 &pcb->tcb_cap);
    if (err) {
        dprintf(6, "procreatE: Failed to retype TCB!\n");
        return err;
    }

    // Configure it
    err = seL4_TCB_Configure(pcb->tcb_cap, pcb->ep_cap, priority,
                             pcb->cspace->root_cnode, seL4_NilData,
                             pcb->seL4_pd, seL4_NilData, SOS_PROCESS_IPC_BUFFER,
                             pcb->ipc_cap);
    if (err) {
        dprintf(6, "procreate: unable to configure new TCB!\n");
        return err;
    }

    return SOS_PROCESS_SUCCESS;
}

static int _load_elf(sos_pcb *pcb, char* app_name, process_create_data *data) {
    // Load the ELF image
    return elf_load(data->pcb, &_start_process, (void *) data);
}

static int _create_stderrout(sos_pcb *pcb) {
    int err = 0;

    // Hacky - redundant open / close on FD 0
    err |= sos_vfs_open(pcb, vfs, "console", O_WRONLY);
    err |= sos_vfs_open(pcb, vfs, "console", O_WRONLY);
    err |= sos_vfs_open(pcb, vfs, "console", O_WRONLY);
    err |= sos_vfs_close(pcb, vfs, 0);
    return err;
}

static void _start_process(int err, seL4_Word pc, void *cookie) {
    process_create_data *data = (process_create_data *) cookie;

    if (err) {
        _fail(err, data);
        return;
    }

    seL4_UserContext context;

    dprintf(5, "\tAttaching stdout and stderr...\n");
    err = _create_stderrout(data->pcb);
    if (err) {
        _fail(err, data);
        return;
    }

    dprintf(5, "\tSetting initial program counter...\n");
    // Set initial seL4 thread context, starting the process
    memset(&context, 0, sizeof(context));
    context.pc = pc;
    context.sp = SOS_PROCESS_STACK_TOP;
    seL4_TCB_WriteRegisters(data->pcb->tcb_cap, 1, 0, 2, &context);
    data->pcb->running = 1;

    // Call back
    _reply(SOS_PROCESS_SUCCESS, data->pcb->pid, data->cb, data->cookie);
    _num_procs++;
    free(data);
    return;
}
